package pl;

import javacard.framework.JCSystem;
import javacard.framework.Util;
import javacard.security.MessageDigest;
import javacard.security.RandomData;

public class EmsaPss {

	// https://www.das-labor.org/trac/browser/microcontroller-2/arm-crypto-lib/testvectors/rsa-pkcs-1v2-1-vec/pss-vect.txt
	// https://github.com/bdauvergne/python-pkcs1/blob/master/tests/data/pss-int.txt
	// http://grepcode.com/file/repo1.maven.org/maven2/org.bouncycastle/bcprov-jdk15on/1.47/org/bouncycastle/crypto/signers/PSSSigner.java
	
	byte[] message = new byte[] { (byte) 0x85, (byte) 0x9E, (byte) 0xEF,
			(byte) 0x2F, (byte) 0xD7, (byte) 0x8A, (byte) 0xCA, (byte) 0x00,
			(byte) 0x30, (byte) 0x8B, (byte) 0xDC, (byte) 0x47, (byte) 0x11,
			(byte) 0x93, (byte) 0xBF, (byte) 0x55, (byte) 0xBF, (byte) 0x9D,
			(byte) 0x78, (byte) 0xDB, (byte) 0x8F, (byte) 0x8A, (byte) 0x67,
			(byte) 0x2B, (byte) 0x48, (byte) 0x46, (byte) 0x34, (byte) 0xF3,
			(byte) 0xC9, (byte) 0xC2, (byte) 0x6E, (byte) 0x64, (byte) 0x78,
			(byte) 0xAE, (byte) 0x10, (byte) 0x26, (byte) 0x0F, (byte) 0xE0,
			(byte) 0xDD, (byte) 0x8C, (byte) 0x08, (byte) 0x2E, (byte) 0x53,
			(byte) 0xA5, (byte) 0x29, (byte) 0x3A, (byte) 0xF2, (byte) 0x17,
			(byte) 0x3C, (byte) 0xD5, (byte) 0x0C, (byte) 0x6D, (byte) 0x5D,
			(byte) 0x35, (byte) 0x4F, (byte) 0xEB, (byte) 0xF7, (byte) 0x8B,
			(byte) 0x26, (byte) 0x02, (byte) 0x1C, (byte) 0x25, (byte) 0xC0,
			(byte) 0x27, (byte) 0x12, (byte) 0xE7, (byte) 0x8C, (byte) 0xD4,
			(byte) 0x69, (byte) 0x4C, (byte) 0x9F, (byte) 0x46, (byte) 0x97,
			(byte) 0x77, (byte) 0xE4, (byte) 0x51, (byte) 0xE7, (byte) 0xF8,
			(byte) 0xE9, (byte) 0xE0, (byte) 0x4C, (byte) 0xD3, (byte) 0x73,
			(byte) 0x9C, (byte) 0x6B, (byte) 0xBF, (byte) 0xED, (byte) 0xAE,
			(byte) 0x48, (byte) 0x7F, (byte) 0xB5, (byte) 0x56, (byte) 0x44,
			(byte) 0xE9, (byte) 0xCA, (byte) 0x74, (byte) 0xFF, (byte) 0x77,
			(byte) 0xA5, (byte) 0x3C, (byte) 0xB7, (byte) 0x29, (byte) 0x80,
			(byte) 0x2F, (byte) 0x6E, (byte) 0xD4, (byte) 0xA5, (byte) 0xFF,
			(byte) 0xA8, (byte) 0xBA, (byte) 0x15, (byte) 0x98, (byte) 0x90,
			(byte) 0xFC };
	
	byte[] hashSHA1 = new byte[] { (byte) 0x37, (byte) 0xB6, (byte) 0x6A,
			(byte) 0xE0, (byte) 0x44, (byte) 0x58, (byte) 0x43, (byte) 0x35,
			(byte) 0x3D, (byte) 0x47, (byte) 0xEC, (byte) 0xB0, (byte) 0xB4,
			(byte) 0xFD, (byte) 0x14, (byte) 0xC1, (byte) 0x10, (byte) 0xE6,
			(byte) 0x2D, (byte) 0x6A };
	
	byte[] saltV = new byte[] { (byte) 0xE3, (byte) 0xB5, (byte) 0xD5,
			(byte) 0xD0, (byte) 0x02, (byte) 0xC1, (byte) 0xBC, (byte) 0xE5,
			(byte) 0x0C, (byte) 0x2B, (byte) 0x65, (byte) 0xEF, (byte) 0x88,
			(byte) 0xA1, (byte) 0x88, (byte) 0xD8, (byte) 0x3B, (byte) 0xCE,
			(byte) 0x7E, (byte) 0x61 };

	public static final byte DEFAULT = 0x00;
	protected MessageDigest md;

	private static final byte[] ZERO_ARRAY = new byte[8];
	private short hLen;
	private RandomData rd;
	short saltLength;

	private byte[] hashResult;

	byte[] salt;
	byte[] result;
	byte[] dbMask;
	byte[] DB;

	byte[] mask;
	byte[] C;
	byte[] hash;

	private short emBits;
	private short sLen;
	short emLen;
	private short counter;
	private short hLen2;

	public EmsaPss(short saltLen, short modLen) {
		
		setDefault();
		
		emBits = (short) 1023; //(modLen - 1);
		
		saltLength = saltLen;
		salt = JCSystem.makeTransientByteArray(saltLength,
				JCSystem.CLEAR_ON_DESELECT); // new byte[saltLen];
		
		C = new byte[(short) 4];
		
		result = new byte[(short) ((short)(emBits + 7) / 8)]; // 128 bytes for 1024 key
		
		result = new byte[(short)128];
	}

	public void setDefault() {
		// default -> sha-1 (20 bytest), mgf1, trailer field 0xBC
		md = MessageDigest.getInstance(MessageDigest.ALG_SHA, false);
		hashResult = JCSystem.makeTransientByteArray(md.getLength(),
				JCSystem.CLEAR_ON_DESELECT);
		hLen = (short) md.getLength();
		rd = RandomData.getInstance(RandomData.ALG_PSEUDO_RANDOM);
		hash = new byte[md.getLength()];
	}

	public void getPssPadding(byte[] hashedMessage) {
		
		hashedMessage = hashSHA1;
			
		//emBits = (short) (md.getLength() * 8 + (short) hashedMessage.length * 8 + 9);
		//emLen = (short) ((short) (emBits + 7) / 8);
		
		rd.generateData(salt, (short) 0, saltLength);
		
		salt = saltV;
		
		emsaPssPadding(hashedMessage, emBits, salt, result);
	}

	private void emsaPssPadding(byte[] hash, short emBits, byte[] salt,
			byte[] outbuff) {
		sLen = (short) salt.length;
		// emLen = (short) ((short)(emBits + 7) / 8);

		if (emBits < (short) (8 * hLen + 8 * sLen + 9)) {
			/* TODO: Exception */
		}

		// hash (00 00 00 00 00 00 00 00 || hash || salt)
		md.reset();
		md.update(ZERO_ARRAY, (short) 0, (short) ZERO_ARRAY.length);
		md.update(hash, (short) 0, (short) hash.length);

		// add salt
		md.doFinal(salt, (short) 0, (short) salt.length, hashResult, (short) 0);
		// OK
		
		// PS len emLen - sLen - hLen - 2 np. 43 B
		
		// DB = PS || 01 || salt.
		DB = null;
		DB = // JCSystem.makeTransientByteArray((short)(emLen - sLen - hLen - 2
		// + 1 + sLen), JCSystem.CLEAR_ON_DESELECT);
		new byte[(short) (emLen - sLen - hLen - 2 + 1 + sLen)];
		// DB len emLen - hLen - 1
		DB[(short) (emLen - sLen - hLen - 2)] = 0x01;

		Util.arrayCopy(salt, (short) 0, DB, (short) (emLen - sLen - hLen - 1),
				sLen);

		// mgf1
		dbMask = null;
		dbMask = mgf1(hashResult, (short) (emLen - hLen - 1));

		for (short i = 0; i < DB.length; i++)
			DB[i] = (byte) (DB[i] ^ dbMask[i]);

		DB[0] &= (0xFF >>> (8 * emLen - emBits));

		// result = new byte[emLen]; //JCSystem.makeTransientByteArray(emLen,
		// JCSystem.CLEAR_ON_RESET); //new byte[emLen];

		Util.arrayCopy(DB, (short) 0, result, (short) 0,
				(short) (emLen - hLen - 1));
		Util.arrayCopy(hashResult, (short) 0, result,
				(short) (emLen - hLen - 1), hLen);

		result[(short) (emLen - 1)] = (byte) 0xBC;

		// return result;
	}

	private byte[] mgf1(byte[] seed, short length) {
		mask = null;
		mask = new byte[length]; // JCSystem.makeTransientByteArray(length,
		// JCSystem.CLEAR_ON_DESELECT); //new
		// byte[length];
		// JCSystem.makeTransientByteArray((short)4,
		// JCSystem.CLEAR_ON_DESELECT); //new byte[4];
		// JCSystem.makeTransientByteArray(md.getLength(),
		// JCSystem.CLEAR_ON_DESELECT);

		counter = 0;
		hLen2 = md.getLength();
		md.reset();

		while (counter < (short) (length / hLen2)) {
			I2OSP(counter, C);
			md.update(seed, (short) 0, (short) seed.length);
			md.doFinal(C, (short) 0, (short) C.length, hash, (short) 0);
			Util.arrayCopy(hash, (short) 0, mask, (short) (counter * hLen2),
					hLen2);
			counter++;
		}

		if ((short) (counter * hLen2) < length) {
			I2OSP(counter, C);
			md.update(seed, (short) 0, (short) seed.length);
			md.doFinal(C, (short) 0, (short) C.length, hash, (short) 0);
			Util.arrayCopy(hash, (short) 0, mask, (short) (counter * hLen2),
					(short) (mask.length - (short) (counter * hLen2)));
		}

		return mask;
	}

	private void I2OSP(short i, byte[] sp) {
		sp[0] = (byte) (i >>> 24);
		sp[1] = (byte) (i >>> 16);
		sp[2] = (byte) (i >>> 8);
		sp[3] = (byte) (i >>> 0);
	}
}
